#!/usr/bin/env python3
import sys
import math
import os, time, random
import json
from pptx import Presentation
from pptx.util import Inches, Pt
from pptx.dml.color import RGBColor
import collections

def parse_outlier_number(output_file):
    fp = open(output_file, "r")

    out_num = 0
    total_num = 0
    for line in fp.readlines():
        line = line.strip()

        total_num += 1
        if int(line) == 0:
            out_num += 1
    fp.close()

    return out_num, total_num

def parse_sess_cost(cost_file):
    cost_list = []

    fp = open(cost_file, "r")
    for line in fp.readlines():
        line = line.strip()

        cost_list.append(line)
    fp.close()

    return cost_list

def parse_abs_fd(fd_file):
    fd_list = []

    fp = open(fd_file, "r")
    for line in fp.readlines():
        line = line.strip()

        fd_list.append(float(line))
    fp.close()

    return sum(fd_list) / len(fd_list)

def fetch_qc_results(rest_dir, qc_reg_dir, sub_list):
    qc_results = collections.OrderedDict()

    for sub_item in os.listdir(rest_dir):
        if sub_item in sub_list:
            qc_results[sub_item] = {}

            sub_dir = os.path.join(rest_dir, sub_item)
            sub_qc_dir = os.path.join(sub_dir, 'qc')
            sub_bold_dir = os.path.join(sub_dir, 'bold')
            sub_bold_mc_dir = os.path.join(sub_bold_dir, 'mc')

            cost_file = 'CBIG_preproc_bbregister_intra_sub_reg.cost'
            full_cost_file = os.path.join(sub_qc_dir, cost_file)
            cost_list = parse_sess_cost(full_cost_file)

            for iter in [1, 2, 3]:
                bld_item = 'bld00' + str(iter)

                qc_results[sub_item][bld_item] = collections.OrderedDict()

                fd_png = sub_item + '_' + bld_item + '_FDRMS0.2_DVARS50.png'
                qc_results[sub_item][bld_item]['fd_img'] = os.path.join(sub_qc_dir, fd_png)

                sub_id = sub_item.split('_')[-1]
                qc_reg_png = 'OHSU_' + sub_id + '_sess_' + str(iter) + '.png'
                full_qc_reg_png = os.path.join(qc_reg_dir, qc_reg_png)

                if os.path.exists(full_qc_reg_png):
                    qc_results[sub_item][bld_item]['qc_reg_img'] = full_qc_reg_png
                else:
                    qc_results[sub_item][bld_item]['qc_reg_img'] = ''

                # meanFD
                fd_file = sub_item + '_' + bld_item + '_rest_skip10_stc_motion_outliers_FDRMS'
                full_fd_file = os.path.join(sub_bold_mc_dir, fd_file)
                mean_fd = parse_abs_fd(full_fd_file)

                qc_results[sub_item][bld_item]['mean_fd'] = round(mean_fd, 2)

                # outlier
                outlier_file = sub_item + '_' + bld_item + '_FDRMS0.2_DVARS50_motion_outliers.txt'
                full_outlier_file = os.path.join(sub_qc_dir, outlier_file)
                outlier_num, total_num = parse_outlier_number(full_outlier_file)
                outlier_rate = float(outlier_num) / total_num

                qc_results[sub_item][bld_item]['outlier_num'] = outlier_num
                qc_results[sub_item][bld_item]['total_num'] = total_num
                qc_results[sub_item][bld_item]['outlier_rates'] = outlier_rate

                if outlier_rate < 0.5:
                    qc_results[sub_item][bld_item]['cost'] = cost_list.pop(0)
                else:
                    qc_results[sub_item][bld_item]['cost'] = 'inf'

    qc_results = [(k,qc_results[k]) for k in sorted(qc_results.keys())]
    order_results = collections.OrderedDict()
    for (sub_name, sub_dict) in qc_results:
        order_results[sub_name] = sub_dict

    return order_results

def parse_sub_list(src_task_file):
    sub_list = []

    fp = open(src_task_file, "r")
    for line in fp.readlines():
        subj_name = line.replace("\n", "")

        sub_list.append(subj_name)

    return sub_list

def generate_pptx(qc_results, trg_pptx_file):
    remian_dict = {}

    prs = Presentation()

    bullet_slide_layout = prs.slide_layouts[1]
    slide = prs.slides.add_slide(bullet_slide_layout)
    shapes = slide.shapes

    title_shape = shapes.title
    body_shape = shapes.placeholders[1]

    title_shape.text = 'Quality control'

    tf = body_shape.text_frame

    p = tf.add_paragraph()
    p.text = '左图头动大小，右图T1、T2配准'
    p.font.size = Pt(20)
    p.level = 1

    p = tf.add_paragraph()
    p.text = 'Outliners大于总时间点数的50% (FD < 0.2)'
    p.font.size = Pt(20)
    p.level = 2

    p = tf.add_paragraph()
    p.text = 'Mean FD > 0.2'
    p.font.size = Pt(20)
    p.level = 2

    p = tf.add_paragraph()
    p.text = 'Intra-subject registration COST > 0.7'
    p.font.size = Pt(20)
    p.level = 2

    for (sub_name, sub_dict) in qc_results.items():
        remian_dict[sub_name] = []

        for (sess_name, sess_dict) in sub_dict.items():
            blank_slide_layout = prs.slide_layouts[6]
            slide = prs.slides.add_slide(blank_slide_layout)

            left = height = Inches(1)
            top = Inches(0.5)
            width = Inches(5)
            txBox = slide.shapes.add_textbox(left, top, width, height)
            tf = txBox.text_frame

            cost = sess_dict['cost']
            outlier = sess_dict['outlier_num']
            total_num = sess_dict['total_num']
            outlier_rates = sess_dict['outlier_rates']
            mean_fd = sess_dict['mean_fd']

            # sub_name
            p = tf.add_paragraph()
            p.text = "Sub ID: {}".format(sub_name)
            p.font.size = Pt(20)
            if cost == 'inf' or float(cost) > 0.7 or outlier_rates > 0.5 or mean_fd > 0.2:
                p.font.color.rgb = RGBColor(0xFF, 0x00, 0x00)
            else:
                remian_dict[sub_name].append(sess_name)

            # sess
            p = tf.add_paragraph()
            p.text = "SessID: {}".format(sess_name + '\n')
            p.font.size = Pt(20)
            if cost == 'inf' or float(cost) > 0.7 or outlier_rates > 0.5 or mean_fd > 0.2:
                p.font.color.rgb = RGBColor(0xFF, 0x00, 0x00)

            # cost, outlier_num, outlier_rate
            p = tf.add_paragraph()
            p.text = "cost: {}, outlier: {}/{}({}), meanFD: {}".format(cost, outlier, total_num, round(outlier_rates, 2), mean_fd)

            img_path = sess_dict['fd_img']
            top = Inches(2.5)
            left = Inches(1)
            height = Inches(4)
            pic = slide.shapes.add_picture(img_path, left, top, height=height)

            if sess_dict['qc_reg_img'] != '':
                img_path = sess_dict['qc_reg_img']
                top = Inches(2.5)
                left = Inches(6)
                height = Inches(3.5)
                pic = slide.shapes.add_picture(img_path, left, top, height=height)
    prs.save(trg_pptx_file)

    return remian_dict

def save_json(remian_dict, trg_path):
    remain_json = json.dumps(remian_dict)

    fp = open(trg_path, 'w')
    fp.write(remain_json)
    fp.close()

if __name__ == '__main__':
    project_dir = '/home/liang/Projects/ASD_QC'
    symbol      = "OHSU"
    sub_type    = "HC"
    list_file   = 'variability_list_' + sub_type + '_' + symbol + '.txt'

    src_task_file = os.path.join(project_dir, os.path.join('Lists/', list_file))
    processed_rest_dir = os.path.join(project_dir, 'DataProcessed/rest/')
    qc_figure_dir = os.path.join(project_dir, 'Figures/QC/{}/{}'.format(symbol, sub_type))
    output_dir = os.path.join(project_dir, 'Record')

    sub_list = parse_sub_list(src_task_file)
    qc_results = fetch_qc_results(processed_rest_dir, qc_figure_dir, sub_list)

    trg_pptx_file = os.path.join(output_dir, sub_type + '_' + symbol + '_reg_qc.pptx')
    remian_dict = generate_pptx(qc_results, trg_pptx_file)
    save_json(remian_dict, 'extract_{}_{}.json'.format(symbol, sub_type))
